library(tidyverse)
library(wesanderson)
library(MetBrewer)
library(sp)
library(sf)
library(classInt)
library(raster)
library(fixest)
library(rgdal)
library(raster)
library(classInt)
library(PNWColors)




        #define some functions
        #########################
                #function for adding transparency to any color	
                add.alpha <- function(col, alpha=1){
                  if(missing(col))
                    stop("Please provide a vector of colours.")
                  apply(sapply(col, col2rgb)/255, 2, 
                        function(x) 
                          rgb(x[1], x[2], x[3], alpha=alpha))  
                }	
                
                get_grid_raster <- function(raster){
                  output <- raster(raster)
                  output[] <- 1:ncell(output)
                  return(output)
                }
                
                identify_non_cells <- function(raster, poly){
                  raster = raster(raster)
                  raster[] <- 1:ncell(raster)
                  r_ext <- extract(raster, poly)
                  all_cells <- 1:ncell(raster)
                  non_cells <- all_cells[all_cells %in% unlist(r_ext) == F]
                  return(non_cells)
                  
                }
        #########################



    #read in maps
      
          map <- read_rds("figdat/world_borders_WB.rds")
            map_bg <- map[map$ADM0_NAME %in% c("India","Pakistan","Nepal","Bhutan","Bangladesh","Afghanistan","Sri Lanka","China","Tajikistan","Myanmar"),]
            map_sa <- map[map$ADM0_NAME %in% c("India","Pakistan","Nepal","Bhutan","Bangladesh","Sri Lanka"),]


            #code section reads in raster data for PM, wealth, and urbanicity and just takes average values for each grid cell and saves.
            #not including raw data here although all is publicly available oneline from those respective authors so could
            #download their data and run this code to generate the average values shown in fig
                
                      ########################################################################################################                
                      ### pm2.5  ###   
              #    
              #             grid <- raster(raster("1_input/data/vandonk/annual_001/V5GL02.HybridPM25.Global.199801-199812.nc"))
              #             grid[] <- 1:ncell(grid)
              #             
              #             
              #             fls <- list.files("1_input/data/vandonk/annual_001/")[13:22][6:10]
              #             
              #             
              #             rdat <- raster::brick(x = grid, nl = length(fls))
              #             
              #             for (i in 1:length(fls)){
              #               
              #               rdat[[i]] <- (raster(paste("data/vandonk/annual_001/", fls[i], sep = "")))
              #               
              #             }
              #             
              #             
              #             r_af <- raster::brick(x = raster::crop(grid,extent(map_sa)), nl = length(fls))
              #             
              #             
              #             
              #             for (i in 1:length(fls)){
              #               r_af[[i]] <- raster::crop(rdat[[i]], extent(map_sa))
              #             }
              #             
              #             pm_sa <- raster::stackApply(x=r_af,indices = 1, fun= mean)
              #             rm(r_af)
              #      
              #             
              #             
              # ### RWI ###
              # 
              #             #loop over countries and read in wealth data
              #             
              #             wfls <- list.files("1_input/data/rwi/")
              #             wfls_cty <- substr(wfls,1,3)
              #             
              #             wfls <- wfls[wfls_cty %in% map_sa$ISO3166_1_]
              #             wfls_cty <- substr(wfls,1,3)
              #             
              #             rwi_list <- list()
              #             
              #             for (i in 1:length(wfls)){
              #               rwi_list[[i]] <- read_csv(paste("1_input/data/rwi/",wfls[i],sep=""))
              #               rwi_list[[i]]$country <- wfls_cty[i]
              #             }      
              #             
              #             
              #             rwi <- data.frame(data.table::rbindlist(rwi_list))
              #             
              #             
              #             rwi$pm <- pm_sa[cellFromXY(pm_sa, as.matrix(rwi[c("longitude","latitude")]))]
              #             rwi <- rwi %>% drop_na(pm, rwi)
              #         dat = rwi
              #         
              #       save(dat,pm_sa, file = "data/bivariate-south-asia-data.Rdata", compress = "gz")    
              
              ####################################################################################################
              ####################################################################################################
              ####################################################################################################
              
                    
                    
                    
      
bipal <- data.frame(bivar_bin = 1:9, col = c("#4885C1","#435786","#3F2949","#89A1C8","#806A8A","#77324C","#CABED0","#BC7C8F","#AE3A4E"))

load("data/south-asia-plot-data.RData")
dat$rwi_bin <- dat$pm_bin <- NA
dat$rwi_bin[dat$rwi < quantile(dat$rwi, .33)] <- 1
dat$rwi_bin[dat$rwi >= quantile(dat$rwi, .33) & dat$rwi <= quantile(dat$rwi, .667)] <- 2
dat$rwi_bin[dat$rwi > quantile(dat$rwi, .667)] <- 3

dat$pm_bin[dat$pm < quantile(dat$pm, .33)] <- 1
dat$pm_bin[dat$pm >= quantile(dat$pm, .33) & dat$pm <= quantile(dat$pm, .667)] <- 2
dat$pm_bin[dat$pm > quantile(dat$pm, .667)] <- 3

dat$bivar_bin <- NA
dat$bivar_bin[dat$rwi_bin == 3 & dat$pm_bin == 1] <- 1
dat$bivar_bin[dat$rwi_bin == 3 & dat$pm_bin == 2] <- 2
dat$bivar_bin[dat$rwi_bin == 3 & dat$pm_bin == 3] <- 3
dat$bivar_bin[dat$rwi_bin == 2 & dat$pm_bin == 1] <- 4
dat$bivar_bin[dat$rwi_bin == 2 & dat$pm_bin == 2] <- 5
dat$bivar_bin[dat$rwi_bin == 2 & dat$pm_bin == 3] <- 6
dat$bivar_bin[dat$rwi_bin == 1 & dat$pm_bin == 1] <- 7
dat$bivar_bin[dat$rwi_bin == 1 & dat$pm_bin == 2] <- 8
dat$bivar_bin[dat$rwi_bin == 1 & dat$pm_bin == 3] <- 9
#adjust cutoffs to better reflect percentiles
dat$bivar_bin[dat$bivar_bin==3 & dat$pm<63] <- 2 
dat$bivar_bin[dat$bivar_bin==8 & dat$pm>57] <- 9
dat$bivar_bin[dat$bivar_bin==8 & dat$pm<41] <- 7
dat$bivar_bin[dat$bivar_bin==5 & dat$pm>58] <- 6
dat$bivar_bin[dat$bivar_bin==1 & dat$pm>39] <- 2

dat <- left_join(dat, bipal)
barplot(table(dat$bivar_bin))





      #prep raster for plotting
      gridsamp <- pm_sa
      gridsamp[]<-NA
      grid10 <- raster::aggregate(gridsamp, fact = 3)
      grid10[cellFromXY(grid10, as.matrix(dat[,c("longitude","latitude")]))] <- dat$bivar_bin
      map_all <- rgeos::gUnaryUnion(map_sa)

      #load cities
      kat <- readOGR("data/city-boundaries/gadm40_NPL_3.shp")
      kat <- kat[kat$NAME_3=="Kathmandu",]
      
      dak <- readOGR("data/city-boundaries/dhaka_administrative_boundaries/dhaka_administrative_boundaries.shp")
      dak = rgeos::gUnaryUnion(dak)
      
      isl <- readOGR("data/city-boundaries/gadm41_PAK_shp/gadm41_PAK_3.shp")
      isl <- isl[isl$NAME_3=="Islamabad",]

     #plot maps for panels e-f       
        
        png(file = "figures/fig3-e-map.png", width = 6000, height = 6000)
        
        par(mar =c(0,0,0,0))

        
        plot(grid10, col = NA, legend = F, bty = "n", box = F,axes=F)
        plot(map_sa, border = NA,col = 'gray90', lwd = 2, add = T)
        

        plot(grid10, col = bipal$col, legend = F, bty = "n", box = F,axes=F, add = T)
        plot(map_sa, border = 'white',col = NA, lwd = 15, add = T)
        plot(map_sa, border = 'black',col = NA, lwd = 5, add = T)
        
        plot(kat, add = T, col = NA, border = 'green', lwd = 2.5)
        plot(isl, add = T, col = NA, border = 'green', lwd = 2.5)
        plot(dak, add = T, col = NA, border = 'green', lwd = 2.5)
        
        dev.off()
        
        


########## City maps ########
      

      
      #prep map data
      pm_k <- raster::crop(pm_sa, extent(kat))
      pm_d <- raster::crop(pm_sa, extent(dak))
      pm_i <- raster::crop(pm_sa, extent(isl))
      
      
      grid_k <- get_grid_raster(pm_k)
      grid_d <- get_grid_raster(pm_d)
      grid_i <- get_grid_raster(pm_i)
      
      noncell_k <- identify_non_cells(grid_k, kat)
      pm_k[noncell_k] <- NA
      
      noncell_d <- identify_non_cells(grid_d, dak)
      pm_d[noncell_d] <- NA
      
      noncell_i <- identify_non_cells(grid_i, isl)
      pm_i[noncell_i] <- NA
  
      
      #combine data for scatters (takes as input raw RWI data which is publicly available for download)
            
            # w_k <- read_csv("data/rwi/NPL_relative_wealth_index.csv") %>% arrange(longitude, latitude)
            # w_loc_k <- SpatialPoints(w_k[,c("longitude","latitude")])
            # crs(w_loc_k)<-crs(kat)
            # overlap <- over(w_loc_k, kat)
            # w_k <- w_k[!is.na(overlap[,1]),]    
            # 
            # 
            # 
            # w_d <- read_csv("data/rwi/BGD_relative_wealth_index.csv") %>% arrange(longitude, latitude)
            # w_loc_d <- SpatialPoints(w_d[,c("longitude","latitude")])
            # crs(w_loc_d)<-crs(dak)
            # overlap <- over(w_loc_d, dak)
            # w_d <- w_d[!is.na(overlap),]  
            # 
            # 
            # 
            # w_i <- read_csv("data/rwi/IND_PAK_relative_wealth_index.csv") %>% arrange(longitude, latitude)
            # w_loc_i <- SpatialPoints(w_i[,c("longitude","latitude")])
            # crs(w_loc_i)<-crs(isl)
            # overlap <- over(w_loc_i, isl)
            # w_i <- w_i[!is.na(overlap[,1]),]  
      
            # save(w_k, w_d, w_i, file = "data/fig3-f-scatter-data.Rdata")
      load("data/fig3-f-scatter-data.Rdata")
      
      #define color pals
      pal_pm <- colorRampPalette(colors = MetBrewer::met.brewer("VanGogh2", 8))(2000) %>% rev()
      pal_wealth <- colorRampPalette(colors = c("red3",MetBrewer::met.brewer("Benedictus", 16)[c(1:6, 8:13)], "navy"))(2000) 


pdf(file = "figures/south-asia-bivar-map-rightpanel-cbf.pdf", width = 8, height = 12)

par(mfrow = c(3,2))

plot(kat, col = NA, border = NA)
plot(pm_k, col = pal_pm, axes = F,add=T, legend = F)
plot(kat, add = T, lwd =2)


plot(kat, col = NA, border = NA)
int <- classIntervals(w_k$rwi, breaks = 100)
col_rio_w <- findColours(int, (pal_wealth))
points(w_k$longitude, w_k$latitude, pch = 15, cex = 3.5, col = col_rio_w)
plot(kat, add = T, lwd =2)




plot(dak, col = NA, border = NA)
plot(pm_d, col = pal_pm, axes = F,add=T, legend = F)
plot(dak, add = T, lwd =2)


plot(dak, col = NA, border = NA)
int <- classIntervals(w_d$rwi, breaks = 100)
col_rio_w <- findColours(int, (pal_wealth))
points(w_d$longitude, w_d$latitude, pch = 15, cex = 5, col = col_rio_w)
plot(dak, add = T, lwd =2)



plot(isl, col = NA, border = NA)
plot(pm_i, col = pal_pm, axes = F,add=T, legend = F)
plot(isl, add = T, lwd =2)






plot(isl, col = NA, border = NA)
int <- classIntervals(w_i$rwi, breaks = 100)
col_rio_w <- findColours(int, (pal_wealth))
points(w_i$longitude, w_i$latitude, pch = 15, cex = 2.8, col = col_rio_w)
plot(isl, add = T, lwd =2)


dev.off()
















##scatters


w_k$pm <- pm_k[cellFromXY(pm_k,as.matrix(w_k[,c("longitude","latitude")]))]
w_d$pm <- pm_d[cellFromXY(pm_d,as.matrix(w_d[,c("longitude","latitude")]))]
w_i$pm <- pm_i[cellFromXY(pm_i,as.matrix(w_i[,c("longitude","latitude")]))]



pdf("figures/south-asia-bivar-map-rightpanel-scatter-withaxes.pdf", width = 3, height = 9)                

par(mfcol =c(3,1))

plot(w_k[,c("rwi","pm")],axes = F, xlab = "", ylab = "", col = NA)
points(w_k[,c("rwi","pm")], pch = 21, bg = add.alpha('gray70', .9), col = add.alpha('white', 1), cex=2, lwd = 0.25)
#      lines(x  =seq(-0.5, 2, .01), y =  predict(loess(pm ~ wealth, data = scatter_pm_wealth_rio, span = 2), newdata = data.frame(wealth = seq(-0.5, 2, .01))), col  = 'blue',lwd=3)
lines(x  =seq(-0.5, 2, .01), y =  predict(lm(pm ~ rwi, data = w_k), newdata = data.frame(rwi = seq(-0.5, 2, .01))), col  = 'blue',lwd=3)
mtext(side = 2, text = "PM2.5",line=3,cex=2)
mtext(side = 1, text = "Wealth",line=3,cex=2)
axis(1, cex.axis = 1.5)
axis(2, las=2, cex.axis = 1.5)


plot(w_d[,c("rwi","pm")],axes = F, xlab = "", ylab = "", col = NA)
points(w_d[,c("rwi","pm")], pch = 21, bg = add.alpha('gray70', .9), col = add.alpha('white', 1), cex=2, lwd = 0.25)
#      lines(x  =seq(-0.5, 2, .01), y =  predict(loess(pm ~ wealth, data = scatter_pm_wealth_rio, span = 2), newdata = data.frame(wealth = seq(-0.5, 2, .01))), col  = 'blue',lwd=3)
lines(x  =seq(-0.5, 2, .01), y =  predict(lm(pm ~ rwi, data = w_d), newdata = data.frame(rwi = seq(-0.5, 2, .01))), col  = 'blue',lwd=3)
mtext(side = 2, text = "PM2.5",line=3,cex=2)
mtext(side = 1, text = "Wealth",line=3,cex=2)
axis(1, cex.axis = 1.5)
axis(2, las=2, cex.axis = 1.5)



plot(w_i[,c("rwi","pm")],axes = F, xlab = "", ylab = "", col = NA)
points(w_i[,c("rwi","pm")], pch = 21, bg = add.alpha('gray70', .9), col = add.alpha('white', 1), cex=2, lwd = 0.25)
#      lines(x  =seq(-0.5, 2, .01), y =  predict(loess(pm ~ wealth, data = scatter_pm_wealth_rio, span = 2), newdata = data.frame(wealth = seq(-0.5, 2, .01))), col  = 'blue',lwd=3)
lines(x  =seq(-0.5, 2, .01), y =  predict(lm(pm ~ rwi, data = w_i), newdata = data.frame(rwi = seq(-0.5, 2, .01))), col  = 'blue',lwd=3)
mtext(side = 2, text = "PM2.5",line=3,cex=2)
mtext(side = 1, text = "Wealth",line=3,cex=2)
axis(1, cex.axis = 1.5)
axis(2, las=2, cex.axis = 1.5)


dev.off()




## Finally generate scatter for panel e scatter insert

mod <- feols(pm ~ rwi | country, data = dat)
dat$pm_scat <- dat$pm

dat$rwi_scat <- dat$rwi
dat$rwi_scat[ dat$rwi_scat>2 ] <- 2



mat <- matrix(nrow =length(0:120)-1, ncol = length(seq(-2, 2, .1)) )
for(i in 1:(nrow(mat) -1)){
  for(j in 1:(ncol(mat)-1 ) ){
    
    mat[i,j]<-sum(
      dat$rwi_scat > seq(-2, 2, .1)[j] & dat$rwi_scat <= seq(-2, 2, .1)[j+1] &
        dat$pm_scat > (0:120)[i] & dat$pm_scat <=  (0:120)[i+1]
      , na.rm =T)
  }
}
mat[mat==0]<-NA




map_y <- data.frame(y_plot=1:(length(0:120)-1)/(length(0:120)-1), y=0:119)
map_x <- data.frame(x_plot = 1:(length(seq(-2, 2, .1)))/length(seq(-2, 2, .1)), x = seq(-2, 2, .1))
new = dat[,c("rwi_scat","pm_scat")]; names(new) = c("x","y")
new[,1] = round(new[,1], 1)
new[,2] = round(new[,2], 0)
map_x = round(map_x, 1)
new <- left_join(new, map_x) %>% left_join(map_y)






png("figures/fig3-e-scatter-insert.png", width = 600, height = 600)                

par(mar = c(4,4,1,0))
plot(flip(raster(mat), direction= "y"), col = pnw_palette("Winter",100)[100:1], legend = F, box = F, axes = F)
axis(1, tick = T, at = seq(0,1,.25), labels = seq(-2, 2,1 ), cex.axis = 2,line=1)
axis(2, tick = T, at = seq(0,1,1/7), labels = seq(0, 70,10 ), las = 2, cex.axis = 2,line=1)

lines(seq(0, 1, .01), predict(lm(y_plot ~ x_plot, new), newdata = data.frame(x_plot = seq(0, 1, .01)) ), col = 'navy',lwd=8)
dev.off()


